package com.lyang.mall.web.framework.redis;

import redis.clients.jedis.Jedis;
import redis.clients.jedis.Transaction;
import redis.clients.jedis.ZParams;

import java.util.List;
import java.util.UUID;

/**
 * redis 限流
 */
public class RedisRateLimiter {
    private static final String BUCKET = "BUCKET";
    private static final String BUCKET_COUNT = "BUCKET_COUNT";
    private static final String BUCKET_MONITOR = "MONITOR";

    /**
     *  从桶中获取令牌
     * @param jedis
     * @param limit
     * @param timeout
     * @return
     */
    public static String acquireTokenFromBucket(Jedis jedis , int limit , long timeout){
        String identifier = UUID.randomUUID().toString();
        long now = System.currentTimeMillis();
        Transaction transaction = jedis.multi();  //开启redis事务

        //删除信号量
        transaction.zremrangeByScore(BUCKET_MONITOR.getBytes(),
                "-inf".getBytes(),String.valueOf(now-timeout).getBytes());
        ZParams params = new ZParams();
        params.weightsByDouble(1.0,1.0);
        transaction.zinterstore(BUCKET,params,BUCKET,BUCKET_MONITOR);

        //计数器自增
        transaction.incr(BUCKET_COUNT);
        List<Object> results = transaction.exec();
        long counter = (Long)results.get(results.size()-1);

        transaction = jedis.multi();
        transaction.zadd(BUCKET_MONITOR,now,identifier);
        transaction.zadd(BUCKET,counter,identifier);
        transaction.zrank(BUCKET,identifier);
        results = transaction.exec();
        //获取排名，判断请求前是否取得了信号量
        long rank = (Long)results.get(results.size()-1);
        if(rank<limit){
            return identifier;
        }else{
            //没有获得信号量，清理之前放入redis中的垃圾数据
            transaction = jedis.multi();
            transaction.zrem(BUCKET_MONITOR,identifier);
            transaction.zrem(BUCKET,identifier);
            transaction.exec();
            return null;
        }
    }
}
